import numpy as np
import matplotlib.pyplot as plt
import pandas as pd

def load_data(excel_file_path):
    return pd.read_excel(excel_file_path)

def extract_column_values(df, column_name):
    return df.iloc[1:, df.columns.get_loc(column_name)].values

def trim_column_values(column_values):
    trim_size = len(column_values) // 5 * 5
    return column_values[:trim_size]

def reshape_values(column_values):
    return column_values.reshape(-1, 5)

def calculate_x_bar_values(column_values_reshaped):
    return np.mean(column_values_reshaped, axis=1)

def calculate_r_values(column_values_reshaped):
    return np.max(column_values_reshaped, axis=1) - np.min(column_values_reshaped, axis=1)

def calculate_control_limits(x_bar_avg, R_avg, A2):
    UCL = x_bar_avg + A2 * R_avg
    LCL = x_bar_avg - A2 * R_avg
    return UCL, LCL

def get_control_values(df, control_standard_column_name, control_column_name):
    df_control = df.iloc[1:, [df.columns.get_loc(control_standard_column_name), df.columns.get_loc(control_column_name)]]
    return df_control.values

def calculate_cp_values(USL_values, LSL_values, sigma):
    return (USL_values - LSL_values) / (6 * sigma)

def calculate_cpk_values(Cp_values, M_values, x_bar_avg, USL_values, LSL_values):
    K_values = 2 * np.abs(M_values - x_bar_avg) / (USL_values - LSL_values)
    return (1 - K_values) * Cp_values


def plot_x_bar_and_r_charts(x_bar_values, R_values, x_bar_avg, UCL, LCL, R_avg, column_name, standard_type, USL_values,
                            LSL_values):
    plt.figure(figsize=(12, 12))

    # X-bar图
    plt.subplot(2, 1, 1)
    plt.plot(x_bar_values, marker='o', linestyle='-')
    plt.axhline(y=x_bar_avg, color='r', linestyle='--', label='中心线')
    plt.axhline(y=UCL, color='g', linestyle='--', label='上限控制限')
    plt.axhline(y=LCL, color='g', linestyle='--', label='下限控制限')

    # 添加规格限
    plt.axhline(y=USL_values[0], color='b', linestyle='--', label='上限规格限')
    plt.axhline(y=LSL_values[0], color='b', linestyle='--', label='下限规格限')

    plt.title(f'X-bar图 ({column_name} - {standard_type})')
    plt.xlabel('样本编号')
    plt.ylabel('平均值')
    plt.legend()
    plt.grid()

    # R图
    plt.subplot(2, 1, 2)
    r_x_values = np.arange(1, len(R_values) + 1)
    plt.plot(r_x_values, R_values, marker='o', linestyle='-')
    plt.axhline(y=R_avg, color='r', linestyle='--', label='中心线')
    plt.title(f'R图 ({column_name} - {standard_type})')
    plt.xlabel('样本编号')
    plt.ylabel('范围 (R)')
    plt.legend()
    plt.grid()

    plt.tight_layout()
    plt.show()

def save_results_to_excel(cp_cpk_df, file_path):
    cp_cpk_df.to_excel(file_path, index=False)

def calculate_and_plot_cp_cpk(excel_file_path, output_file_path):
    plt.rcParams['font.sans-serif'] = ['SimHei']
    plt.rcParams['axes.unicode_minus'] = False

    df = load_data(excel_file_path)

    columns_to_extract = ["MaterialTopDS", "MaterialTopOS", "MaterialBottomDS", "MaterialBottomOS", "TotalDS", "TotalOS"]

    Cp_values_list = []
    Cpk_values_list = []

    for column_name in columns_to_extract:
        column_values = extract_column_values(df, column_name)
        column_values = trim_column_values(column_values)
        column_values_reshaped = reshape_values(column_values)
        x_bar_values = calculate_x_bar_values(column_values_reshaped)
        R_values = calculate_r_values(column_values_reshaped)

        A2 = 0.729
        d2 = 2.059

        x_bar_avg = np.mean(x_bar_values)
        R_avg = np.mean(R_values)
        UCL, LCL = calculate_control_limits(x_bar_avg, R_avg, A2)

        control_column_name = column_name + "Control"
        control_standard_column_name = column_name + "ControlStandard"

        control_values = get_control_values(df, control_standard_column_name, control_column_name)

        USL_values = control_values[:, 0] + control_values[:, 1]
        LSL_values = control_values[:, 0] - control_values[:, 1]

        standard_type = df['standardType'][1]

        plot_x_bar_and_r_charts(x_bar_values, R_values, x_bar_avg, UCL, LCL, R_avg, column_name, standard_type, USL_values, LSL_values)

        sigma = R_avg / d2
        Cp_values = calculate_cp_values(USL_values, LSL_values, sigma)

        M_values = control_values[:, 0]
        Cpk_values = calculate_cpk_values(Cp_values, M_values, x_bar_avg, USL_values, LSL_values)

        Cp_values_list.append(Cp_values)
        Cpk_values_list.append(Cpk_values)

    cp_cpk_df = pd.DataFrame({'Column': columns_to_extract, 'Cp': Cp_values_list, 'Cpk': Cpk_values_list})
    save_results_to_excel(cp_cpk_df, output_file_path)

# 示例用法:
if __name__ == '__main__':
    excel_file_path = "yySPCtest.xlsx"
    output_file_path = "Cp_Cpk_results.xlsx"
    calculate_and_plot_cp_cpk(excel_file_path, output_file_path)
